#include <ros/ros.h>
#include <nav_msgs/Odometry.h>
#include <geometry_msgs/Twist.h>
#include <tf/transform_broadcaster.h>
#include <sensor_msgs/Imu.h>
#include <sensor_msgs/LaserScan.h>

class RobotOdom
{
public:
  RobotOdom() :
    x_(0.0), y_(0.0), theta_(0.0),
    vx_(0.0), vy_(0.0), vth_(0.0)
  {
    // Initialize ROS node
    ros::NodeHandle nh_;

    // Subscribe to cmd_vel topic
    cmd_vel_sub_ = nh_.subscribe("cmd_vel", 10, &RobotOdom::cmdVelCallback, this);

    // Advertise odom topic
    odom_pub_ = nh_.advertise<nav_msgs::Odometry>("encode_odom", 50);

    // Initialize tf broadcaster
    tf_broadcaster_ = std::make_shared<tf::TransformBroadcaster>();

    // Initialize IMU and Laser subscribers
    imu_sub_ = nh_.subscribe("imu/data", 10, &RobotOdom::imuCallback, this);
    // laser_sub_ = nh_.subscribe("scan", 10, &RobotOdom::laserCallback, this);

    // Set the update rate
    // update_rate_ = 50.0; // Hz
    //timer_ = nh_.createTimer(ros::Duration(1.0 / update_rate_), &RobotOdom::update, this);
  }

private:
  void cmdVelCallback(const geometry_msgs::Twist::ConstPtr& msg)
  {
    vx_ = msg->linear.x;
    vy_ = msg->linear.y;
    vth_ = msg->angular.z;

    //update();
  }

  void imuCallback(const sensor_msgs::Imu::ConstPtr& msg)
  {

    if (msg->orientation_covariance[0] < 0)
        return ;
    // Use IMU data to improve orientation estimation
    // For simplicity, we assume the IMU provides a reliable orientation
    tf::Quaternion q(msg->orientation.x, msg->orientation.y, msg->orientation.z, msg->orientation.w);
    tf::Matrix3x3 m(q);
    double roll, pitch, yaw;
    m.getRPY(roll, pitch, yaw);
    theta_ = yaw;
  }

  void laserCallback(const sensor_msgs::LaserScan::ConstPtr& msg)
  {
    // Use laser data to improve position estimation
    // For simplicity, we assume the laser data is used to correct position drift
    // This can be done using SLAM algorithms or other techniques
  }

  void update()
  {
    // Update position based on velocity
    double dt = 1.0 / update_rate_;
    double delta_x = (vx_ * cos(theta_) - vy_ * sin(theta_)) * dt;
    double delta_y = (vx_ * sin(theta_) + vy_ * cos(theta_)) * dt;
    double delta_theta = vth_ * dt;

    x_ += delta_x;
    y_ += delta_y;
    theta_ += delta_theta;

    // Publish odometry message
    nav_msgs::Odometry odom;
    odom.header.stamp = ros::Time::now();
    odom.header.frame_id = "odom";
    odom.child_frame_id = "base_footprint";

    // Set the position
    odom.pose.pose.position.x = x_;
    odom.pose.pose.position.y = y_;
    odom.pose.pose.position.z = 0.0;
    odom.pose.pose.orientation = tf::createQuaternionMsgFromYaw(theta_);

    // Set the velocity
    odom.twist.twist.linear.x = vx_;
    odom.twist.twist.linear.y = vy_;
    odom.twist.twist.angular.z = vth_;

    // Publish the message
    odom_pub_.publish(odom);

    // Broadcast the transform
    // geometry_msgs::TransformStamped odom_trans;
    // odom_trans.header.stamp = odom.header.stamp;
    // odom_trans.header.frame_id = "odom";
    // odom_trans.child_frame_id = "base_footprint";

    // odom_trans.transform.translation.x = x_;
    // odom_trans.transform.translation.y = y_;
    // odom_trans.transform.translation.z = 0.0;
    // odom_trans.transform.rotation = odom.pose.pose.orientation;

    // tf_broadcaster_->sendTransform(odom_trans);
  }

  ros::NodeHandle nh_;
  ros::Subscriber cmd_vel_sub_;
  ros::Subscriber imu_sub_;
  ros::Subscriber laser_sub_;
  ros::Publisher odom_pub_;
  std::shared_ptr<tf::TransformBroadcaster> tf_broadcaster_;
  ros::Timer timer_;

  double x_, y_, theta_;
  double vx_, vy_, vth_;
  double update_rate_;
};

int main(int argc, char** argv)
{
  ros::init(argc, argv, "robot_odom");
  RobotOdom robot_odom;
  ros::spin();
  return 0;
}
